# coding: utf-8
import tensorflow as tf
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense


def build_model():
    model = tf.keras.Sequential([
        # 输入层
        tf.keras.layers.Reshape((28, 28, 1), input_shape=(784,)),

        # 第一卷积层
        tf.keras.layers.Conv2D(32, (5, 5), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),

        # 第二卷积层
        tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),

        # 全连接层
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(1024, activation='relu'),
        tf.keras.layers.Dropout(0.5),#随机丢弃50%的神经元，防止过拟合

        # 输出层
        tf.keras.layers.Dense(10, activation='softmax')#激活函数使用softmax
    ])
    return model


if __name__ == '__main__':
    # 加载数据
    (train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

    # 数据预处理
    train_images = train_images.reshape(-1, 784).astype('float32') / 255.0#展开成列向量，并归一化处理
    test_images = test_images.reshape(-1, 784).astype('float32') / 255.0
    train_labels = tf.keras.utils.to_categorical(train_labels, 10)#对标签独热编码
    test_labels = tf.keras.utils.to_categorical(test_labels, 10)
    # 构建模型
    model = build_model()

    # 编译模型
    model.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])#Adam（学习率 0.0001），用于更新模型权重，使用categorical_crossentropy交叉熵计算损失
    # 训练模型
    history=model.fit(train_images, train_labels,
              batch_size=50,
              epochs=10,
              validation_split=0.1)

    # 保存完整模型
    model.save('mnist_cnn_model.h5')
    np.save('cnn_train_history.npy',history.history)
    # 保存处理后的测试数据用于后续评估
    np.savez('test_data.npz', images=test_images, labels=test_labels)